#pragma once

#include <string>
#include <vector>
#include <memory>
#include <map>
#include <libflexio/flexio.h>
#include <doca_dpa.h>
#define DEVICE_TYPE_NUM 3
#define MAX_LAYER_NUM 10
#define MAX_DPA_THREAD_NUM 256
#define EMBEDDING_DIM 128

#define USING_BRPC
// #define USING_KV
// #define USING_SEQUENCE
// 必须叫这个名
extern "C"
{
    flexio_func_t sequence_preprocess;
    flexio_func_t dpa_execute_unit_init;
    flexio_func_t test_preprocess;
    flexio_func_t threadcontrol_init;
    flexio_func_t set_addr;
    flexio_func_t get_max_thread_time;
    doca_dpa_func_t hello_world;
    doca_dpa_func_t docasdkprocess;
}
extern struct flexio_app *feature_device;

struct OperatorNode;
struct MemoryNode;

struct MemoryNode {
    std::string name;
    int id = -1;
    int feature_start_id = 0;
    int feature_end_id = 1;
    int batch_start_id = 0;
    int batch_end_id = 1;
    int element_per_feature = 1;
    //TODO 建图的时候把内存分配了，一开始先posix_align分配一整块内存，然后这里再按照operator去计算分配。

    std::string dtype = "int64";
    size_t total_size = 0;
    void* storage_ptr = nullptr; // 每个memory node分配到的内存起始地址。

    std::vector<std::shared_ptr<OperatorNode>> prev_opt_nodes;  // 原 producers
    std::vector<std::shared_ptr<OperatorNode>> next_opt_nodes;  // 原 consumers

    std::vector<std::shared_ptr<MemoryNode>> prev_memory_nodes;
    std::vector<std::shared_ptr<MemoryNode>> next_memory_nodes;
};

typedef enum {
    DEVICE_CPU,  // 默认从0开始（DEVICE_CPU = 0）
    DEVICE_DPA,  // 1
    DEVICE_ARM   // 2
} device_type_t;

typedef enum {
    FLEXIO_SDK,
    DOCA_SDK,
    TEST_SDK
} sdk_type_t;

typedef enum {
    NORMAL_STAGE,
    SEND_STAGE,
    RECV_STAGE
} stage_type_t;

struct OperatorNode {
    std::string op_type;
    std::string input_name; //输入特征名，特征名字+#第几个
    std::string output_name;
    device_type_t device_type; //分配到哪种device上执行。
    int id = -1;
    int thread_num = 0; // 用几个thread_num去做这个operatornode
    std::map<std::string, std::string> params;

    std::vector<std::shared_ptr<MemoryNode>> prev_mem_nodes;  // 原 input_nodes
    std::vector<std::shared_ptr<MemoryNode>> next_mem_nodes;  // 原 output_nodes

    std::vector<std::shared_ptr<OperatorNode>> prev_ops;
    std::vector<std::shared_ptr<OperatorNode>> next_ops;
};

struct ExecuteNode{
    int op_type;
    void * input_addr;
    void * output_addr;
};

class CommonUtils{
public:
    /* alloc a posix memory */
    static void * prepare_host_memory(size_t size) {
        void * ptr = nullptr;
        int ret = posix_memalign(&ptr, 64, (size + 63) & (~63));
        if (ret != 0) {
            printf ("error malloc host memory\n");
        }
        return ptr;
    }

    static int calc_ele(int ele_num, std::string op_type) {
        if (op_type == "ngram") { //TODO: 用上operator_node 里的 params
            int32_t actual_n_of_gram = std::min(4, ele_num);
            int32_t final_length = (ele_num - actual_n_of_gram + 1) * actual_n_of_gram;
            return final_length;
        }
        if (op_type == "firstx") {
            return 1;
        }
        if (op_type == "onehot") {
            return ele_num * 8;
        }
        return ele_num;
    }

    static void memory_dtype_assign(std::shared_ptr<OperatorNode> operation_node) {
        for (auto mem_node : operation_node->prev_mem_nodes) {
            if (operation_node->op_type == "logit" || 
                operation_node->op_type == "boxcox" ||
                operation_node->op_type == "bucketize" ||
                operation_node->op_type == "onehot") {
                    mem_node->dtype = "double";
                }
        }

        for (auto mem_node : operation_node->next_mem_nodes) {
            if (operation_node->op_type == "logit" || 
                operation_node->op_type == "boxcox" ||
                operation_node->op_type == "onehot") {
                    mem_node->dtype = "double";
                }
            
            if (operation_node->op_type == "embedding_fetch") {
                mem_node->dtype = "float";
            }
        }
    }

    static bool tmp_mem_node(std::shared_ptr<MemoryNode> mem_node_a, std::shared_ptr<MemoryNode> mem_node_b){
        return mem_node_a->id < mem_node_b->id;
    }

    template<typename T>
    static void mem_init(T * ptr, size_t size) {
        size_t num = size / sizeof(T);
        for (size_t i = 0; i < num; ++i) {
            ptr[i] = static_cast<T>(i + 10);
        }
    }
};